import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
import torchvision.models as models
from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights

def FastRCNN_model(num_classes):
    # 加载 Faster R-CNN 预训练模型
    weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT  # 使用默认的最新权重
    model = models.detection.fasterrcnn_resnet50_fpn(weights=weights)
    # 获取输入特征数
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # 定义新的头部为检测器，带有所需数量的类
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    return model

